from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime
import json
import uuid
from pathlib import Path

from utils.logger import Logger
from utils.exceptions import AgentError, ToolError
from utils.constants import AGENT_STATUSES, TASK_STATUSES
from states.agent_state import AgentState
from states.task_state import TaskState
from .document_manager import DocumentManager
from .segment_analyzer import SegmentAnalyzer
from auxiliary_agents.compliance import check_compliance
from distributed.circuit_breaker import circuit_breaker
from distributed.fallback import fallback_handler

logger = get_logger("doc_segment_agent")


@dataclass
class DocSegmentRequest:
    """文档分段检索请求数据结构"""
    query: str  # 查询文本
    user_id: str  # 用户ID
    task_id: str  # 任务ID
    doc_type: str = "auto"  # 文档类型: pdf, docx, txt, auto
    segment_type: str = "chapter"  # 分段类型: chapter, page, paragraph, section
    metadata_filter: Optional[Dict[str, Any]] = None  # 元数据过滤条件
    top_k: int = 5  # 返回前K个最相关段落


@dataclass
class DocSegmentResponse:
    """文档分段检索响应数据结构"""
    task_id: str
    segments: List[Dict[str, Any]]  # 检索到的段落列表
    status: str  # 参考 TASK_STATUSES
    confidence: float  # 总体置信度
    metadata: Dict[str, Any]  # 文档元数据信息
    error_msg: Optional[str] = None
    timestamp: str = datetime.utcnow().isoformat() + "Z"


class DocSegmentAgent:
    """文档分段检索Agent：处理文档的元数据提取、分段和智能检索"""

    def __init__(self):
        self.agent_id = "doc_segment_agent"
        self.agent_state = AgentState(
            agent_id=self.agent_id,
            status=AGENT_STATUSES.RUNNING,
            last_heartbeat=datetime.utcnow()
        )
        self.document_manager = DocumentManager()
        self.segment_analyzer = SegmentAnalyzer()
        self.similarity_threshold = 0.7  # 相似度阈值
        logger.info(f"DocSegmentAgent initialized: agent_id={self.agent_id}")

    def _update_agent_state(self, status: str):
        """更新Agent状态"""
        self.agent_state.status = status
        self.agent_state.last_heartbeat = datetime.utcnow()
        logger.debug(f"Agent state updated: {asdict(self.agent_state)}")

    def _validate_request(self, request: Dict[str, Any]) -> DocSegmentRequest:
        """验证请求参数合法性"""
        required_fields = ["query", "user_id", "task_id"]
        for field in required_fields:
            if field not in request or not request[field]:
                raise AgentError(f"请求缺少必填字段: {field}")

        return DocSegmentRequest(
            query=str(request["query"]).strip(),
            user_id=str(request["user_id"]).strip(),
            task_id=str(request["task_id"]).strip(),
            doc_type=request.get("doc_type", "auto"),
            segment_type=request.get("segment_type", "chapter"),
            metadata_filter=request.get("metadata_filter", {}),
            top_k=int(request.get("top_k", 5))
        )

    def _extract_document_metadata(self, file_path: str, doc_type: str) -> Dict[str, Any]:
        """提取文档元数据"""
        try:
            metadata = {
                "file_path": file_path,
                "doc_type": doc_type,
                "file_size": Path(file_path).stat().st_size,
                "created_time": datetime.fromtimestamp(Path(file_path).stat().st_ctime).isoformat(),
                "modified_time": datetime.fromtimestamp(Path(file_path).stat().st_mtime).isoformat(),
                "page_count": 0,
                "chapter_count": 0,
                "title": "",
                "author": "",
                "description": ""
            }

            # 根据文档类型提取特定元数据
            if doc_type == "pdf":
                metadata.update(self.document_manager.extract_pdf_metadata(file_path))
            elif doc_type == "docx":
                metadata.update(self.document_manager.extract_docx_metadata(file_path))
            elif doc_type == "txt":
                metadata.update(self.document_manager.extract_txt_metadata(file_path))

            return metadata
        except Exception as e:
            logger.error(f"提取文档元数据失败: {file_path}, error={str(e)}")
            return {}

    def _segment_document(self, file_path: str, segment_type: str) -> List[Dict[str, Any]]:
        """对文档进行分段处理"""
        try:
            segments = []
            
            if segment_type == "chapter":
                segments = self.segment_analyzer.extract_chapters(file_path)
            elif segment_type == "page":
                segments = self.segment_analyzer.extract_pages(file_path)
            elif segment_type == "paragraph":
                segments = self.segment_analyzer.extract_paragraphs(file_path)
            elif segment_type == "section":
                segments = self.segment_analyzer.extract_sections(file_path)
            else:
                # 智能分段
                segments = self.segment_analyzer.intelligent_segment(file_path)

            # 为每个段落添加元数据
            for i, segment in enumerate(segments):
                segment["segment_id"] = f"seg_{i}_{uuid.uuid4().hex[:8]}"
                segment["index"] = i
                segment["timestamp"] = datetime.utcnow().isoformat() + "Z"

            return segments
        except Exception as e:
            logger.error(f"文档分段失败: {file_path}, error={str(e)}")
            return []

    @circuit_breaker(fallback=fallback_handler, threshold=5, timeout=30)
    def process_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """处理文档分段检索任务（核心入口）"""
        task_id = task.get("task_id", "unknown_task")
        logger.info(f"开始处理文档分段检索任务: task_id={task_id}, query={task.get('query')}")

        try:
            self._update_agent_state(AGENT_STATUSES.BUSY)

            # 1. 验证请求
            request = self._validate_request(task)

            # 2. 合规性检查
            compliance_result = check_compliance(request.query)
            if not compliance_result["allowed"]:
                logger.warning(f"任务不合规: task_id={task_id}, reason={compliance_result['reason']}")
                response = DocSegmentResponse(
                    task_id=task_id,
                    segments=[],
                    status=TASK_STATUSES.REJECTED,
                    confidence=0.0,
                    metadata={},
                    error_msg=f"查询不合规: {compliance_result['reason']}"
                )
                return asdict(response)

            # 3. 文档加载和元数据提取
            if "file_path" in task:
                file_path = task["file_path"]
                doc_type = request.doc_type if request.doc_type != "auto" else self.document_manager.detect_doc_type(file_path)
                
                # 提取文档元数据
                metadata = self._extract_document_metadata(file_path, doc_type)
                
                # 文档分段
                segments = self._segment_document(file_path, request.segment_type)
                
                # 构建向量索引
                self.document_manager.build_segment_index(task_id, segments)
                
                # 4. 段落检索
                relevant_segments = self.segment_analyzer.search_relevant_segments(
                    query=request.query,
                    segments=segments,
                    top_k=request.top_k,
                    similarity_threshold=self.similarity_threshold
                )
                
                # 5. 计算总体置信度
                confidence = sum(seg.get("similarity", 0) for seg in relevant_segments) / len(relevant_segments) if relevant_segments else 0.0
                
                # 6. 构造响应
                response = DocSegmentResponse(
                    task_id=task_id,
                    segments=relevant_segments,
                    status=TASK_STATUSES.COMPLETED if relevant_segments else TASK_STATUSES.NO_RESULT,
                    confidence=confidence,
                    metadata=metadata
                )
                
                logger.info(f"任务处理完成: task_id={task_id}, status={response.status}, segments={len(relevant_segments)}")
                return asdict(response)
            else:
                # 基于已有索引的检索
                relevant_segments = self.document_manager.search_from_index(
                    query=request.query,
                    top_k=request.top_k,
                    metadata_filter=request.metadata_filter
                )
                
                confidence = sum(seg.get("similarity", 0) for seg in relevant_segments) / len(relevant_segments) if relevant_segments else 0.0
                
                response = DocSegmentResponse(
                    task_id=task_id,
                    segments=relevant_segments,
                    status=TASK_STATUSES.COMPLETED if relevant_segments else TASK_STATUSES.NO_RESULT,
                    confidence=confidence,
                    metadata={"search_type": "index_search"}
                )
                
                logger.info(f"索引检索完成: task_id={task_id}, status={response.status}, segments={len(relevant_segments)}")
                return asdict(response)

        except ToolError as e:
            logger.error(f"工具调用失败: task_id={task_id}, error={str(e)}")
            response = DocSegmentResponse(
                task_id=task_id,
                segments=[],
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                metadata={},
                error_msg=f"工具调用失败: {str(e)}"
            )
            return asdict(response)

        except AgentError as e:
            logger.error(f"Agent处理失败: task_id={task_id}, error={str(e)}")
            response = DocSegmentResponse(
                task_id=task_id,
                segments=[],
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                metadata={},
                error_msg=f"任务处理失败: {str(e)}"
            )
            return asdict(response)

        except Exception as e:
            logger.error(f"未知错误: task_id={task_id}, error={str(e)}", exc_info=True)
            response = DocSegmentResponse(
                task_id=task_id,
                segments=[],
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                metadata={},
                error_msg=f"系统错误: {str(e)}"
            )
            return asdict(response)

        finally:
            self._update_agent_state(AGENT_STATUSES.RUNNING)

    def get_agent_state(self) -> Dict[str, Any]:
        """获取Agent当前状态"""
        return asdict(self.agent_state)


# 单例Agent实例
doc_segment_agent = DocSegmentAgent()

# 测试入口
if __name__ == "__main__":
    test_task = {
        "task_id": "test_doc_001",
        "query": "第三章讲了什么内容？",
        "user_id": "test_user",
        "file_path": "/tmp/test.pdf",
        "segment_type": "chapter"
    }

    result = doc_segment_agent.process_task(test_task)
    print("测试结果:", result)